# from multiprocessing import process as task
import ssh, time, thread, Report, pickle, gtk
from subprocess import call, Popen, PIPE
from ssh import Connection
from grid import ScrollGrid

class TestRunner():
	def __init__(self):	
		try:
			self.tests = pickle.load(open('tests.p', 'rb'))
		except IOError:
			self.tests = []
		# connected instance pointers
		self.clients = []
		self.clientConnections = []

		# debug variables
		self.verbose = 1
		self.connect = 1

		# global defines:
		self.privateKey = '/home/eli/id_rsa-gsg-keypair'
#		self.clientType = 'ami-71fd1a18'     #Default ubuntu server image
		self.clientType = 'ami-f037d099'     #custom grinder image
		self.scriptLocation = '/home/eli/DISTRIBUTE/scripts'
		self.grinderLocation = '/home/eli/DISTRIBUTE/grinder3-2'
		self.environmentLocation = '/home/eli/DISTRIBUTE/environment'
		
		self.startGui()
		
		
	def startGui(self):
		print "start GUI"

	def startPool(self, target):
		if self.verbose: print('need %s instances.' % target)
		self.findClients()
		existing = len(self.clients)
		needed = target - existing
		if self.verbose: print('found %s existing instances, starting %s.' % (existing, needed))
		for x in range(existing, target):
			self.spawnClient(self.clientType, x)
		if needed > 0:
			if self.verbose: print('waiting for new servers to boot')
			time.sleep(60)
			self.findClients()
		if self.verbose: print('done finding clients')

	def spawnClient(self, ami, x):
		if self.verbose: print('initializing server %s.' % (x + 1))
		cmd = ['ec2-run-instances', ami, '-k', 'gsg-keypair']
		Popen(cmd, stdout=PIPE)

		
	def findClients(self):
		if self.verbose: print('checking instances...')
		cmd = ['ec2-describe-instances']
		outputData = Popen(cmd, stdout=PIPE)
		self.extractClients(outputData)
		if self.verbose: print('found %s clients, connected to %s.' % (len(self.clients), 
									       len(self.clientConnections)))

	def extractClients(self, outputData):
		self.clients, self.clientConnections = [], []
		while True:
			o = outputData.stdout.readline()
			if o == '' and outputData.poll() != None: break
			if 'amazonaws.com' in o:
				self.addClient(o)

	def addClient(self, hostString):
		host = self.getHost(hostString)
		self.clients.append(host)
		if self.verbose: print('found instance: %s' % host)
		try:
			self.clientConnections.append(self.connectTo(host))
			if self.verbose: print('connected to: %s' % host)
		except: print('error connecting to client %s.' % host)

	def connectTo(self, host):
		return ssh.Connection(host, private_key=self.privateKey, username='root')
	   	
	def prepAllClients(self):
		for x in self.clients:
			self.prepClient(x)
			
	def prepClient(self, client):
		if self.verbose: print('starting copy to: %s' % client)
		self.copyFiles(client)
		#self.restartSystem(client)
		#self.installSoftware(client)
	
	def copyFiles(self, client):
		cmd1 = ['scp', '-i', self.privateKey, self.environmentLocation, 'root@%s:/etc' % client]
		cmd2 = ['scp', '-r', '-i', self.privateKey, self.grinderLocation, 'root@%s:/home' % client]
		Popen(cmd1)
		Popen(cmd2)

	def restartAll(self):
		if self.verbose: print('restarting all servers, please wait a few seconds before reconnecting.')
		for client in self.clientConnections: self.restartSystem(client)
	
	def restartSystem(self, client):
		client.execute('reboot')
	
	def installSoftware(self, client):  # B0rkEN
		s = ssh.Connection(client, private_key=self.privateKey, username='root')
		s.execute('sudo apt-get update')
		s.execute('sudo apt-get -y install sun-java6-jre')
		time.sleep(2)
		s.execute('yes')
	
	def generateConnectStrings(self):
		for x in self.clients: 
			print('ssh -i %s root@%s' % (self.privateKey, x))
		
	def startGrinders(self):
		for x in self.clientConnections: 
			try:
				thread.start_new_thread(self.runGrinder, (x,))
			except:
				print('error starting grinder thread: %s' % x)
				x.close()

	def runGrinder(self, client):
		if self.verbose: print('starting a grinder...')
		output = client.execute('java net.grinder.Grinder /home/scripts/grinder.properties')
		if self.verbose: print(output)
		client.close()
			
	def addBox(self):
		if self.verbose: print('adding annother box to pool')
		self.spawnClient(self.clientType)
		self.findClients()
			
	def refreshScripts(self):
		for x in self.clients:
			self.sendScript(x)
		time.sleep(10 * len(self.clients))
		if self.verbose: print('finished updating scripts')
			
	def sendScript(self, client):
		if self.verbose: print('updating instance scripts: %s' % client)
		cmd = ['scp', '-r', '-i', self.privateKey, self.scriptLocation, 'root@%s:/home' % client]
		Popen(cmd)	

	def closeConnections(self):
		for client in self.clientConnections:
			client.close()
		
	def killAll(self):
		if self.verbose: print('killing all clients')
		for client in self.clientConnections:
			self.terminateClient(client)
		
	def terminateClient(self, client):
		try: client.execute('/sbin/shutdown -h now')
		except: print('failed to shutdown instance %s.' % client)
		
	def getHost(self, strIn):
		return strIn.split('amazonaws.com')[0].split('\t')[-1] + 'amazonaws.com'
		
#TESTING
tr = TestRunner()
#tr.startPool(5)
#tr.refreshScripts()
#tr.startGrinders()

#KILL TESTS
#time.sleep(30)
#print('closing connections after 30 seconds')
#tr.killAll()
#tr.closeConnections()

